// SPDX-License-Identifier: GPL-2.0

/**
 * Copyright (C) 2025 Huawei Technologies Co., Ltd.
 * 
 * This program is free software; you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation; either version 2 of the License, or
 * (at your option) any later version.
 * 
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 * 
 * You should have received a copy of the GNU General Public License along
 * with this program; if not, see <https://www.gnu.org/licenses/>.
 */

#include <iostream>
#include <unordered_map>
#include "cmanuapi.h"

#include <boost/program_options.hpp>

namespace po = boost::program_options;

static hw_info_ioctl_args_t hw_info;

struct Configuration
{
    struct Common
    {
        int start_cpu = -1;
        int end_cpu = -1;
        enum class Mode
        {
            unspecified,
            l2_restrict
        } mode = Mode::unspecified;

    } common;

    struct l2Restrict
    {
        unsigned long dmask = 0;
        unsigned long imask = 0;
    } l2_restrict;

};

static std::ostream& operator<<(std::ostream& stream, Configuration::Common::Mode m)
{
    switch (m) {
    case Configuration::Common::Mode::unspecified:
        stream << "SoC configuration features";
        break;
    case Configuration::Common::Mode::l2_restrict:
        stream << "L2 restrict";
        break;
    default:
        assert(0);
    }

    return stream;
}

static void validate(boost::any& v,
    std::vector<std::string> const& values,
    Configuration::Common::Mode*,
    int)
{
    using namespace boost::program_options;

    validators::check_first_occurrence(v);

    const std::string& s = validators::get_single_string(values);

    Configuration::Common::Mode result = Configuration::Common::Mode::unspecified;
    if (s == "l2-restrict") {
        result = Configuration::Common::Mode::l2_restrict;
    } else {
        throw validation_error(validation_error::invalid_option_value);
    }

    v = boost::any(Configuration::Common::Mode(result));
}

static bool parse_hex_input(std::string str, unsigned long *val)
{
    try {
        size_t pos;
        *val = std::stoul(str, &pos, 16);
        if (pos != str.length()) return false;
        return true;
    }
    catch (const std::invalid_argument&) {
        std::cerr << "error:input is invalid_argument" << std::endl;
        return false;
    }
    catch (const std::out_of_range&) {
        std::cerr << "error: input out_of_range" << std::endl;
        return false;
    }
}

static void print_generic_error_msg(Configuration::Common::Mode mode, int cpu_id)
{
    std::cerr << "Current CPU " << cpu_id << " on Mode " << mode << ": operation failed with error - \'" <<
                strerror(errno) << "\'"<< std::endl;
}

static void validate_cpu(po::variables_map parsedOptions, Configuration conf)
{
    int min_cpu_id = -1;

    if (parsedOptions.count("start_cpu") && parsedOptions.count("end_cpu")) {
        min_cpu_id = 0;
    } else if (conf.common.mode == Configuration::Common::Mode::l2_restrict &&
        !parsedOptions.count("reset") && !parsedOptions.count("d-mask") && !parsedOptions.count("i-mask")) {
        std::cerr << "Error: the option \'--start_cpu,-s --end_cpu,-e\' is required for \'" <<
            conf.common.mode << " get config\' mode, but missing" << std::endl;
        exit(1);
    }
    if (conf.common.start_cpu < min_cpu_id || conf.common.start_cpu >= sysconf(_SC_NPROCESSORS_CONF)) {
        std::cerr << "Error: start_cpu ID "<< conf.common.start_cpu << " is out of boundaries" << std::endl;
        exit(1);
    }
    if (conf.common.end_cpu < min_cpu_id || conf.common.end_cpu >= sysconf(_SC_NPROCESSORS_CONF)) {
        std::cerr << "Error: end_cpu ID "<< conf.common.end_cpu << " is out of boundaries" << std::endl;
        exit(1);
    }
    if(conf.common.start_cpu > conf.common.end_cpu) {
        std::cerr << "Error: start_cpu ID: "<< conf.common.start_cpu << " is larger than end_cpu ID:" << conf.common.end_cpu <<  std::endl;
        exit(1);
    }
}

static std::vector<ssize_t>
handle_features_reset(Configuration conf)
{
    size_t all_cpu_num = sysconf(_SC_NPROCESSORS_CONF);
    std::vector<ssize_t> reset_info(all_cpu_num, -1);

    switch (conf.common.mode) {
    case Configuration::Common::Mode::l2_restrict:
        for(auto cpu_id = conf.common.start_cpu; cpu_id <= conf.common.end_cpu; cpu_id++) {
            reset_info[cpu_id] = cman_uapi_reset_l2_partition(cpu_id);
        }
        break;
    default:
        reset_info.assign(all_cpu_num, 1);
    }
    return reset_info;
}

struct CpuDealInfo{
    unsigned long dmask = 0;
    unsigned long imask = 0;
    ssize_t ret = -1;
};

static std::pair<std::string, std::vector<CpuDealInfo>>
l2_restrict_mode_exec(po::variables_map parsedOptions, Configuration conf, std::string l2_dmask, std::string l2_imask)
{
    std::string op;
    size_t all_cpu_num = sysconf(_SC_NPROCESSORS_CONF);
    std::vector<CpuDealInfo> vec_info(all_cpu_num);

    std::cout << "d-mask " << parsedOptions.count("d-mask") << ", i-mask " << parsedOptions.count("i-mask") << std::endl;

    if (parsedOptions.count("d-mask") ^ parsedOptions.count("i-mask")) {
        std::cerr << "Error: Data mask (--d-mask) and Instruction mask (--i-mask) should either "
                     "be specified together or both be unspecified" << std::endl;
        exit(1);
    } else if (parsedOptions.count("d-mask") && parsedOptions.count("i-mask")) {
        if (!parse_hex_input(l2_dmask, &conf.l2_restrict.dmask)) {
            std::cerr << "Error: --d-mask contains non-hexadecimal characters" << std::endl;
            exit(1);
        }
        if (!parse_hex_input(l2_imask, &conf.l2_restrict.imask)) {
            std::cerr << "Error: --i-mask contains non-hexadecimal characters" << std::endl;
            exit(1);
        }
        op = "set";
        for(auto cpu_id = conf.common.start_cpu; cpu_id <= conf.common.end_cpu; cpu_id++) {
            vec_info[cpu_id].ret = cman_uapi_l2_partition_set(cpu_id, conf.l2_restrict.dmask, conf.l2_restrict.imask);
        }
    } else {
        op = "get";
        for(auto cpu_id = conf.common.start_cpu; cpu_id <= conf.common.end_cpu; cpu_id++) {
            vec_info[cpu_id].ret = cman_uapi_l2_partition_get(cpu_id, &vec_info[cpu_id].dmask, &vec_info[cpu_id].imask);
        }
    }
    return std::make_pair(op, vec_info);
}

int main(int argc, char* argv[])
{
    Configuration conf;
    std::string l2_dmask, l2_imask;
    ssize_t ret = -1;
    char str[300];

    ret = cman_uapi_init();
    if (ret < 0) {
        std::cout << "cman_uapi_init() returned: " << ret <<
                    " errno: " << strerror(errno) << std::endl;
        exit(1);
    }
    hw_info = cman_uapi_hw_info();

    po::options_description commonOptions("Common options");

    commonOptions.add_options()
        ("help,h", "Print this help")
        ("start_cpu,s", po::value<int>(&conf.common.start_cpu), "The start number of CPU to configure. This is a required option\n")
        ("end_cpu,e", po::value<int>(&conf.common.end_cpu), "The end number of CPU to configure. This is a required option\n")
        ("mode,m", po::value<Configuration::Common::Mode>(&conf.common.mode), "benchmark (=feature) mode"
            "\nSoC configuration features:"
            "\n    l2-restrict \tL2 instruction-data partitioning feature. Limits effective L2 cache size for instructions and/or data accesses.")
        ("reset", "Resets SoC configuration feature specified by \'--mode,-m\' for CPU specified by \'--start_cpu,-s --end_cpu,-e\' . \n")
    ;

    po::options_description l2RestrictOptions("L2 restrict options");

    sprintf(str,
            "Data mask to be set for L2 cache restriction (in hex: e.g. 0xFF). "
            "Each bit setting means effective L2 data cache size reduction by %lld kilobytes."
            "\n    \tLimits: 0x0 - 0x%llx",
            hw_info.l2_segment_size / 1024,
            (1ull << (hw_info.l2_capacity / hw_info.l2_segment_size)) - 1);
    l2RestrictOptions.add_options()
        ("d-mask", po::value<std::string>(&l2_dmask), str)
    ;
    sprintf(str,
            "Instruction mask to be set for L2 cache restriction (in hex: e.g. 0xFF). "
            "Each bit setting means effective L2 instruction cache size reduction by %lld kilobytes."
            "\n    \tLimits: 0x0 - 0x%llx",
            hw_info.l2_segment_size / 1024,
            (1ull << (hw_info.l2_capacity / hw_info.l2_segment_size)) - 1);
    l2RestrictOptions.add_options()
        ("i-mask", po::value<std::string>(&l2_imask), str)
    ;

    po::options_description options;

    options.add(commonOptions).add(l2RestrictOptions);

    po::variables_map parsedOptions;
    try{
        po::store(po::parse_command_line(argc, argv, options), parsedOptions);
        if (parsedOptions.count("help")) {
            std::cerr << "Usage: ./devcman [-h] [-m <mode>] [-s <start_cpu>] [-e <end_cpu>] [--reset] [feature options]" << std::endl;
            std::cerr << commonOptions;
            std::cerr << "Note:\na)command-line options unrelated to the specified <feature> are ignored.\n" << std::endl;
            std::cerr << l2RestrictOptions;
            std::cerr << "Note:\na) \'--d-mask\' & \'--i-mask\' should be specified together. Specifying only one will result in error. "
                "If none of these options are specified, then the tool will operate in `get L2 restrict config` mode.\n"
                "b) \'--start_cpu,-s --end_cpu,-e\' should be specified in `get L2 restrict config` mode.\n"
                "c) It is not allowed to mask entire L2 cache, at least 2 bits in mask should be unset.\n" << std::endl;
            exit(0);
        }
        po::notify(parsedOptions);
    } catch(const po::error& e) {
        std::cerr << "error: " << e.what() << std::endl;
        std::cerr << commonOptions << std::endl;
        std::cerr << l2RestrictOptions << std::endl;
        return 1;
    }

    validate_cpu(parsedOptions, conf);

    size_t all_cpu_num = sysconf(_SC_NPROCESSORS_CONF);
    std::vector<ssize_t> reset_info(all_cpu_num, ret);
    if (parsedOptions.count("reset")) {
        bool cpu_reset_success = true;
        reset_info = handle_features_reset(conf);
        for(auto cpu_id = conf.common.start_cpu; cpu_id <= conf.common.end_cpu; cpu_id++) {
            if (reset_info[cpu_id] < 0) {
                cpu_reset_success = false;
                std::cerr <<"CPU: " << cpu_id <<  "set Mode " << conf.common.mode << " reset: operation failed with error - \'" <<
                    strerror(errno) << "\'"<< std::endl;
            }
            if (reset_info[cpu_id] > 0) {
                cpu_reset_success = false;
                std::cerr << "CPU: " << cpu_id << " Reset operation is not supported for the specified \'--mode,-m\'." << std::endl;
            }
        }
        if(cpu_reset_success) {
            std::cout <<"Start_CPU: " << conf.common.start_cpu  << " to End_CPU: " << conf.common.end_cpu <<
            " set Mode " << conf.common.mode << " reset: success!" << std::endl;
        }
        return 0;
    }

    switch (conf.common.mode) {
    case Configuration::Common::Mode::l2_restrict:
        {
        auto pair_info = l2_restrict_mode_exec(parsedOptions, conf, l2_dmask, l2_imask);
        bool cpu_set_success = true;
        for(auto cpu_id = conf.common.start_cpu; cpu_id <= conf.common.end_cpu; cpu_id++) {
            if(pair_info.second[cpu_id].ret < 0) {
                print_generic_error_msg(conf.common.mode, cpu_id);
                cpu_set_success = false;
            }
        }

        if(cpu_set_success) {
            if (pair_info.first == "set") {
                std::cout << "Start_CPU: " << conf.common.start_cpu  << " to End_CPU: " << conf.common.end_cpu <<
                " set Mode " << conf.common.mode << " " << pair_info.first <<
                             ": success!" << std::endl;
            } else {
                for(auto cpu_id = conf.common.start_cpu; cpu_id <= conf.common.end_cpu; cpu_id++) {
                    std::cout << "Mode " << conf.common.mode << " " << pair_info.first <<
                    ": success! CPU: " << cpu_id << "; D-mask: 0x" << std::hex <<
                    pair_info.second[cpu_id].dmask << ", I-mask: 0x" << pair_info.second[cpu_id].imask << std::dec <<
                    ";" << std::endl;
                }
            }
        }
        break;
        }
    default:
        std::cerr << "Error: the option \'--mode,-m\' is required but missing" << std::endl;
    }

    return 0;
}
